{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import mindspore\n",
    "from mindspore.dataset import text, GeneratorDataset, transforms\n",
    "from mindspore import nn\n",
    "\n",
    "from mindnlp import load_dataset\n",
    "from mindnlp.transforms import PadTransform, GPTTokenizer\n",
    "\n",
    "from mindnlp.engine import Trainer, Evaluator\n",
    "from mindnlp.engine.callbacks import CheckpointCallback, BestModelCallback\n",
    "from mindnlp.metrics import Accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "imdb_ds = load_dataset('imdb', split=['train', 'test'])\n",
    "imdb_train = imdb_ds['train']\n",
    "imdb_test = imdb_ds['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "imdb_train.get_dataset_size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def process_dataset(dataset, tokenizer, max_seq_len=256, batch_size=32, shuffle=False):\n",
    "    def pad_sample(text):\n",
    "        if len(text) + 2 >= max_seq_len:\n",
    "            return np.concatenate(\n",
    "                [np.array([tokenizer.bos_token_id]), text[: max_seq_len-2], np.array([tokenizer.eos_token_id])]\n",
    "            )\n",
    "        else:\n",
    "            pad_len = max_seq_len - len(text) - 2\n",
    "            return np.concatenate( \n",
    "                [np.array([tokenizer.bos_token_id]), text,\n",
    "                 np.array([tokenizer.eos_token_id]),\n",
    "                 np.array([tokenizer.pad_token_id] * pad_len)]\n",
    "            )\n",
    "\n",
    "    column_names = [\"text\", \"label\"]\n",
    "    rename_columns = [\"input_ids\", \"label\"]\n",
    "\n",
    "    if shuffle:\n",
    "        dataset = dataset.shuffle(batch_size)\n",
    "\n",
    "    # map dataset\n",
    "    dataset = dataset.map(operations=[tokenizer, pad_sample], input_columns=\"text\")\n",
    "    dataset = dataset.map(operations=transforms.TypeCast(mindspore.int32), input_columns=\"label\")\n",
    "    # rename dataset\n",
    "    dataset = dataset.rename(input_columns=column_names, output_columns=rename_columns)\n",
    "    # batch dataset\n",
    "    dataset = dataset.batch(batch_size)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tokenizer\n",
    "gpt_tokenizer = GPTTokenizer.from_pretrained('openai-gpt')\n",
    "\n",
    "# add sepcial token: <PAD>\n",
    "special_tokens_dict = {\n",
    "    \"bos_token\": \"<bos>\",\n",
    "    \"eos_token\": \"<eos>\",\n",
    "    \"pad_token\": \"<pad>\",\n",
    "}\n",
    "num_added_toks = gpt_tokenizer.add_special_tokens(special_tokens_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split train dataset into train and valid datasets\n",
    "imdb_train, imdb_val = imdb_train.split([0.7, 0.3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = process_dataset(imdb_train, gpt_tokenizer, shuffle=True)\n",
    "dataset_val = process_dataset(imdb_val, gpt_tokenizer)\n",
    "dataset_test = process_dataset(imdb_test, gpt_tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindnlp.models import GPTForSequenceClassification\n",
    "from mindnlp._legacy.amp import auto_mixed_precision\n",
    "\n",
    "# set bert config and define parameters for training\n",
    "model = GPTForSequenceClassification.from_pretrained('openai-gpt', num_labels=2)\n",
    "model.pad_token_id = gpt_tokenizer.pad_token_id\n",
    "model.resize_token_embeddings(model.config.vocab_size + 3)\n",
    "model = auto_mixed_precision(model, 'O1')\n",
    "\n",
    "loss = nn.CrossEntropyLoss()\n",
    "optimizer = nn.Adam(model.trainable_params(), learning_rate=2e-5)\n",
    "\n",
    "metric = Accuracy()\n",
    "\n",
    "# define callbacks to save checkpoints\n",
    "ckpoint_cb = CheckpointCallback(save_path='checkpoint', ckpt_name='gpt_imdb_finetune', epochs=1, keep_checkpoint_max=2)\n",
    "best_model_cb = BestModelCallback(save_path='checkpoint', ckpt_name='gpt_imdb_finetune_best', auto_load=True)\n",
    "\n",
    "trainer = Trainer(network=model, train_dataset=dataset_train,\n",
    "                  eval_dataset=dataset_train, metrics=metric,\n",
    "                  epochs=3, loss_fn=loss, optimizer=optimizer, callbacks=[ckpoint_cb, best_model_cb],\n",
    "                  jit=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# start training\n",
    "trainer.run(tgt_columns=\"label\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "evaluator = Evaluator(network=model, eval_dataset=dataset_test, metrics=metric)\n",
    "evaluator.run(tgt_columns=\"label\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
